"""
书法文本识别模块
负责识别书法作品中的文字内容，并判断其出处
"""

import os
import logging
import traceback
from typing import Dict, List, Any, Optional, Tuple
import numpy as np
import cv2
from paddleocr import PaddleOCR
import openai
import requests
import json
from PIL import Image

# 配置日志
logger = logging.getLogger(__name__)

# 配置OpenAI API (如果使用)
try:
    openai.api_key = os.environ.get("OPENAI_API_KEY", "")
except Exception as e:
    logger.warning(f"无法初始化OpenAI API: {str(e)}")

class TextRecognizer:
    """
    书法文本识别器
    负责识别书法中的文字，并判断出处
    """
    
    def __init__(self):
        """初始化识别器"""
        try:
            logger.info("初始化书法文本识别器...")
            # 尝试初始化PaddleOCR
            self.ocr = PaddleOCR(
                use_angle_cls=True,  # 使用角度分类器
                lang="ch",  # 中文模型
                use_gpu=False,  # 不使用GPU
                show_log=False  # 不显示日志
            )
            logger.info("文本识别器初始化完成")
        except Exception as e:
            logger.error(f"初始化OCR引擎失败: {str(e)}")
            logger.error(traceback.format_exc())
            self.ocr = None
    
    def recognize_text(self, image_path: str) -> Dict[str, Any]:
        """
        识别图片中的文字，并判断出处
        
        Args:
            image_path: 图片路径
            
        Returns:
            包含识别结果的字典
        """
        if self.ocr is None:
            return {"error": "OCR引擎未初始化"}
        
        try:
            # 加载图像
            image = cv2.imread(image_path)
            if image is None:
                return {"error": "无法加载图像"}
            
            # 预处理图像以提高OCR效果
            preprocessed_image = self._preprocess_image(image)
            
            # 使用OCR识别文字
            result = self.ocr.ocr(preprocessed_image, cls=True)
            
            # 提取识别到的文本
            recognized_text = self._extract_text_from_result(result)
            logger.info(f"识别到的文本: {recognized_text}")
            
            # 判断文本出处
            source_info = self._identify_text_source(recognized_text)
            
            # 组织返回结果
            return {
                "success": True,
                "recognized_text": recognized_text,
                "source_info": source_info,
                "ocr_details": result  # 可选，返回详细的OCR结果
            }
            
        except Exception as e:
            logger.error(f"文本识别过程中出错: {str(e)}")
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e)
            }
    
    def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """
        预处理图像以提高OCR效果
        
        Args:
            image: 原始图像
            
        Returns:
            处理后的图像
        """
        try:
            # 转为灰度图
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            
            # 自适应二值化，处理不同光照条件
            binary = cv2.adaptiveThreshold(
                gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                cv2.THRESH_BINARY_INV, 11, 2
            )
            
            # 去噪
            denoised = cv2.fastNlMeansDenoising(binary, None, 10, 7, 21)
            
            # 返回处理后的图像
            return denoised
        except Exception as e:
            logger.warning(f"图像预处理失败，使用原图: {str(e)}")
            return image
    
    def _extract_text_from_result(self, result) -> str:
        """
        从OCR结果中提取文本
        
        Args:
            result: OCR识别结果
            
        Returns:
            提取的文本字符串
        """
        if not result or not result[0]:
            return ""
        
        # 提取文本并按位置排序
        text_blocks = []
        for line in result[0]:
            if len(line) >= 2:  # 确保数据结构符合预期
                text = line[1][0]  # 文本内容
                confidence = line[1][1]  # 置信度
                box = line[0]  # 边界框
                
                # 计算中心点y坐标，用于排序
                center_y = sum(point[1] for point in box) / 4
                
                text_blocks.append({
                    "text": text,
                    "confidence": confidence,
                    "center_y": center_y,
                    "box": box
                })
        
        # 根据y坐标排序，模拟从上到下的阅读顺序
        text_blocks.sort(key=lambda x: x["center_y"])
        
        # 组合文本
        full_text = "".join(block["text"] for block in text_blocks)
        return full_text
    
    def _identify_text_source(self, text: str) -> Dict[str, Any]:
        """
        判断文本的出处
        
        Args:
            text: 识别到的文本
            
        Returns:
            出处信息字典
        """
        if not text or len(text) < 5:  # 文本太短，无法判断
            return {
                "identified": False,
                "reason": "文本过短或未识别到内容"
            }
        
        try:
            # 首先尝试使用本地知识库判断
            result = self._local_source_identification(text)
            if result["identified"]:
                return result
            
            # 如果本地识别失败，尝试使用OpenAI API
            if openai.api_key:
                return self._openai_source_identification(text)
                
            # 如果以上方法都失败
            return {
                "identified": False,
                "reason": "无法识别出处"
            }
        except Exception as e:
            logger.error(f"判断文本出处时出错: {str(e)}")
            return {
                "identified": False,
                "error": str(e)
            }
    
    def _local_source_identification(self, text: str) -> Dict[str, Any]:
        """
        使用本地知识库判断文本出处
        
        Args:
            text: 识别到的文本
            
        Returns:
            出处信息字典
        """
        # 这里可以实现简单的关键词匹配或模式识别
        # 例如常见的古诗文判断
        
        # 当前仅实现一个非常简单的示例
        common_poems = {
            "春眠不觉晓": {
                "title": "春晓",
                "author": "孟浩然",
                "dynasty": "唐",
                "full_text": "春眠不觉晓，处处闻啼鸟。夜来风雨声，花落知多少。"
            },
            "床前明月光": {
                "title": "静夜思",
                "author": "李白",
                "dynasty": "唐",
                "full_text": "床前明月光，疑是地上霜。举头望明月，低头思故乡。"
            },
            "白日依山尽": {
                "title": "送杜少府之任蜀州",
                "author": "王勃",
                "dynasty": "唐",
                "full_text": "城阙辅三秦，风烟望五津。与君离别意，同是宦游人。海内存知己，天涯若比邻。无为在歧路，儿女共沾巾。"
            },
            "锄禾日当午": {
                "title": "悯农",
                "author": "李绅",
                "dynasty": "唐",
                "full_text": "锄禾日当午，汗滴禾下土。谁知盘中餐，粒粒皆辛苦。"
            },
            "香九龄": {
                "title": "三字经",
                "author": "王应麟",
                "dynasty": "宋",
                "full_text": "三字经是中国古代蒙学经典之一..."
            }
        }
        
        # 检查文本是否包含常见古诗的关键句
        for key, info in common_poems.items():
            if key in text:
                return {
                    "identified": True,
                    "title": info["title"],
                    "author": info["author"],
                    "dynasty": info["dynasty"],
                    "full_text": info["full_text"],
                    "match_confidence": "高",
                    "method": "关键句匹配"
                }
        
        return {
            "identified": False,
            "reason": "本地库中未找到匹配"
        }
    
    def _openai_source_identification(self, text: str) -> Dict[str, Any]:
        """
        使用OpenAI API判断文本出处
        
        Args:
            text: 识别到的文本
            
        Returns:
            出处信息字典
        """
        if not openai.api_key:
            return {
                "identified": False,
                "reason": "未配置OpenAI API密钥"
            }
        
        try:
            # 构建prompt
            prompt = f"""
            以下是从中国书法作品中识别出的文字，请判断其出处：
            
            "{text}"
            
            请提供以下信息（如果可以确定）：
            1. 这段文字出自哪部作品？
            2. 作者是谁？
            3. 朝代
            4. 完整原文
            5. 匹配的可信度（高/中/低）
            
            如果无法确定，请说明原因。请以JSON格式返回，包含以下字段：identified(bool), title, author, dynasty, full_text, confidence, reason（如果identified为false）
            """
            
            # 调用API
            response = openai.Completion.create(
                engine="gpt-3.5-turbo-instruct",  # 根据可用模型调整
                prompt=prompt,
                max_tokens=500,
                temperature=0.3
            )
            
            # 解析结果
            result_text = response.choices[0].text.strip()
            
            try:
                # 尝试解析JSON
                result = json.loads(result_text)
                result["method"] = "OpenAI API"
                return result
            except json.JSONDecodeError:
                # 如果不是有效的JSON，进行文本解析
                logger.warning(f"OpenAI返回的结果不是有效的JSON: {result_text}")
                
                # 简单文本解析
                lines = result_text.split('\n')
                parsed_result = {
                    "identified": "无法确定" not in result_text and "不确定" not in result_text,
                    "method": "OpenAI API (文本解析)",
                    "raw_response": result_text
                }
                
                for line in lines:
                    if ":" in line:
                        key, value = line.split(":", 1)
                        key = key.strip().lower()
                        value = value.strip()
                        
                        if "作品" in key or "标题" in key or "出处" in key:
                            parsed_result["title"] = value
                        elif "作者" in key:
                            parsed_result["author"] = value
                        elif "朝代" in key:
                            parsed_result["dynasty"] = value
                        elif "原文" in key or "全文" in key:
                            parsed_result["full_text"] = value
                        elif "可信度" in key or "匹配度" in key or "置信度" in key:
                            parsed_result["match_confidence"] = value
                
                return parsed_result
                
        except Exception as e:
            logger.error(f"OpenAI API调用失败: {str(e)}")
            return {
                "identified": False,
                "reason": f"API调用失败: {str(e)}"
            }

# 初始化全局识别器实例
text_recognizer = None

def get_text_recognizer():
    """获取文本识别器实例"""
    global text_recognizer
    if text_recognizer is None:
        text_recognizer = TextRecognizer()
    return text_recognizer 